# encoding: utf-8

"""
    token的计算
"""

from typing import Tuple, List

import tiktoken

from conf.config import config, logger


def get_tokens_from_string(string_raw: str, model_name: str) -> Tuple[int, List]:
    """
    获取字符串对应的token列表及token数量
    :param string_raw: 要转化为token的原始字符串
    :param model_name: 使用的openai模型名
    :return: token数量及token列表
    """

    encoding = tiktoken.encoding_for_model(model_name)
    token_list = encoding.encode(string_raw)
    token_count = len(token_list)
    return token_count, token_list


def main():
    string_raw = "1+2的结果"

    token_count, token_list = get_tokens_from_string(
        string_raw=string_raw,
        model_name=config.get("open_ai_model")
    )
    logger.info(f"原字符数量：{len(string_raw)}")
    logger.info(f"token数量：{token_count}")
    logger.info(f"token列表：{token_list}")


if __name__ == '__main__':
    main()
